from torch import nn
from torchvision import transforms
from torch import utils
import torch
from torchsummary import summary
from dataset import GetDataLoader

class LeNet(nn.Module):
    PTH_PATH = 'model_files/lenet-best-model.pth' # 训练好的模型参数保存的文件路径
    IMG_CHANNELS = GetDataLoader().IMG_CHANNELS  # 输入的图片的通道数
    IMG_CLASSES = len(GetDataLoader().IMG_CLASSES) # 数据集图片的分类数量
    def __init_params(self):
        ## 初始化参数
        for m in self.modules():
            # 卷积层参数的初始化，w和b参数
            if isinstance(m, nn.Conv2d):
                ## 对于卷积，使用凯明初始化（对w初始化）
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                ## 对b参数进行初始
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            # 全连接成参数的初始化
            elif isinstance(m, nn.Linear):
                ## 默认使用正态分布
                nn.init.normal_(m.weight, 0, 0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
    def __init__(self):
        super(LeNet,self).__init__()
        ## 初始参数
        self.__init_params()
        ## 网络结构
        self.sequent = nn.Sequential(
            # 一次卷积（参数6*5*5=150），一次激活函数
            nn.Conv2d(in_channels=self.IMG_CHANNELS,out_channels=6,kernel_size=5,padding=2),
            nn.Sigmoid(),#无参数

            # 平均池化，用于压缩特征（无参数）
            nn.AvgPool2d(kernel_size=2,stride=2),

            # 一次卷积（参数5*5*16=400），一次激活函数
            nn.Conv2d(in_channels=6,out_channels=16,kernel_size=5),
            nn.Sigmoid(),#无参数

            # 平均池化，压缩特征（无参数）
            nn.AvgPool2d(kernel_size=2,stride=2),

            # 展平（无参数）
            nn.Flatten(),

            # 全连接层：参数：400*120
            nn.Linear(in_features=400,out_features=120),
            # 全连接层：参数：120*84
            nn.Linear(in_features=120,out_features=84),
            # 全连接层：参数：84*10
            nn.Linear(in_features=84,out_features=self.IMG_CLASSES)
        )

    def forward(self,input):
        '''
        :param input: 输入的图片是(1,28,28)  1是通道数
        :return:
        '''
        output = self.sequent(input)
        return output


if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = LeNet()
    model = model.to(device)
    print(summary(model, (3, 28, 28)))